﻿Imports System.Net.Sockets
Imports System.Threading
Imports System.Net
Imports System.Runtime.InteropServices

Public Class CantContinueException
    Inherits Exception

    Public Sub New(ByVal pMessage As String, Optional ByVal pDetail As String = "")
        MyBase.New(pMessage)
        Trace.WriteLine("Can't continue: " & pMessage & vbNewLine & pDetail)
    End Sub
End Class

''' <summary>
''' Represents an IMAP command in progress
''' </summary>
''' <remarks></remarks>
Public Interface ICommand
    ''' <summary>
    ''' See if the command will accept a line of input.
    ''' </summary>
    ''' <param name="pLine">The line of input to check for acceptance</param>
    ''' <returns>True if the line of input is accepted</returns>
    ''' <remarks></remarks>
    Function Process(ByVal pLine As String) As Boolean

    ''' <summary>
    ''' Is this command completed
    ''' </summary>
    ''' <value></value>
    ''' <returns></returns>
    ''' <remarks></remarks>
    ReadOnly Property IsCompleted() As Boolean

    ''' <summary>
    ''' Send this command over the wire.
    ''' </summary>
    ''' <remarks></remarks>
    Sub Send()
End Interface

''' <summary>
''' We want increasing indexes to use to identify our commands.
''' </summary>
''' <remarks></remarks>
Public Class CommandIndex
    Private Shared m_Lock As New Object
    Private Shared m_Index As Integer = 0

    ''' <summary>
    ''' Get the next available index for a command prefix.
    ''' </summary>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Public Shared Function GetNextIndex() As Integer
        SyncLock m_Lock
            Dim _ind As Integer = m_Index
            m_Index += 1
            Return _ind
        End SyncLock
    End Function
End Class

''' <summary>
''' Base functionality for commands.
''' </summary>
''' <remarks></remarks>
Public MustInherit Class BaseCommand
    Implements ICommand

    Protected m_Prefix As String
    Protected m_Client As ImapClient

    Public Sub New(ByVal pClient As ImapClient)
        m_Client = pClient
        m_Prefix = CommandIndex.GetNextIndex.ToString("IMAP0000")
    End Sub

    Protected m_IsCompleted As Boolean

    Public Overridable ReadOnly Property IsCompleted() As Boolean Implements ICommand.IsCompleted
        Get
            Return m_IsCompleted
        End Get
    End Property

    Public Function Process(ByVal pLine As String) As Boolean Implements ICommand.Process
        ' a * means untagged input
        If pLine.StartsWith("*") Then
            Return ProcessUntagged(pLine)
        End If

        ' a + means a continuation
        If pLine.StartsWith("+") Then
            Return ProcessContinuation(pLine)
        End If

        ' if a line starts with out prefix, then it is directed at us.
        If pLine.StartsWith(m_Prefix) Then
            pLine = pLine.Substring(m_Prefix.Length + 1, pLine.Length - (m_Prefix.Length + 1))
            If pLine.StartsWith("OK") Then
                System.Diagnostics.Trace.WriteLine("processing ok")
                Return ProcessOK(pLine)
            ElseIf pLine.StartsWith("BAD") Then
                System.Diagnostics.Trace.WriteLine("processing bad")
                Return ProcessBAD(pLine)
            ElseIf pLine.StartsWith("NO") Then
                System.Diagnostics.Trace.WriteLine("processing no")
                Return ProcessNO(pLine)
            Else
                System.Diagnostics.Trace.WriteLine("processing response")
                Return ProcessResponse(pLine)
            End If
        Else
            Return False
        End If
    End Function

    ''' <summary>
    ''' Default handling for an OK being returned
    ''' </summary>
    ''' <param name="pLine"></param>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Public Overridable Function ProcessOK(ByVal pLine As String) As Boolean
        m_IsCompleted = True
        Return True
    End Function

    ''' <summary>
    ''' Default handling for a BAD being returned
    ''' </summary>
    ''' <param name="pLine"></param>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Public Overridable Function ProcessBAD(ByVal pLine As String) As Boolean
        m_IsCompleted = True
        Return True
    End Function

    ''' <summary>
    ''' Default handling for a NO being returned
    ''' </summary>
    ''' <param name="pLine"></param>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Public Overridable Function ProcessNO(ByVal pLine As String) As Boolean
        m_IsCompleted = True
        Return True
    End Function

    ''' <summary>
    ''' Default handling for a response being returned.
    ''' </summary>
    ''' <param name="pLine"></param>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Public Overridable Function ProcessResponse(ByVal pLine As String) As Boolean
        Return False
    End Function

    ''' <summary>
    ''' Default handling for an untagged response being returned
    ''' </summary>
    ''' <param name="pLine"></param>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Public Overridable Function ProcessUntagged(ByVal pLine As String) As Boolean
        Return False
    End Function

    ''' <summary>
    ''' Default handling for a continuation being returned.
    ''' </summary>
    ''' <param name="pLine"></param>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Public Overridable Function ProcessContinuation(ByVal pLine As String) As Boolean
        Return False
    End Function

    ''' <summary>
    ''' Derived classes must implement how to send their output.
    ''' </summary>
    ''' <remarks></remarks>
    Public MustOverride Sub Send() Implements ICommand.Send

    ''' <summary>
    ''' Trace the command we are running
    ''' </summary>
    ''' <param name="pCmd"></param>
    ''' <remarks></remarks>
    Public Overridable Sub Trace(ByVal pCmd As String)
        System.Diagnostics.Trace.WriteLine("Sending: " & pCmd)
    End Sub

    ''' <summary>
    ''' Tell the client to send, and attach our prefix.
    ''' </summary>
    ''' <param name="pCmd"></param>
    ''' <remarks></remarks>
    Public Sub SendCore(ByVal pCmd As String)
        Dim _cmd As String = m_Prefix & " " & pCmd & vbCrLf
        Trace(_cmd)
        m_Client.Send(_cmd)
    End Sub
End Class

''' <summary>
''' Command to check the capabilities of the IMAP server.
''' </summary>
''' <remarks></remarks>
Public Class CapabilitiesCommand
    Inherits BaseCommand

    Private Const CMD_NAME As String = "CAPABILITY"
    Private m_CanIdle As Boolean = False

    Public Sub New(ByVal pClient As ImapClient)
        MyBase.New(pClient)
    End Sub

    Public Overrides Function ProcessUntagged(ByVal pLine As String) As Boolean
        'our application depends on the availability of the IDLE command. Check if we can do it.
        If pLine.Contains(CMD_NAME) Then
            If pLine.Contains("IDLE") Then
                m_CanIdle = True
            End If
        End If
        Return True
    End Function

    Public Overrides Sub Send()
        Dim _cmd As String = CMD_NAME
        SendCore(_cmd)
    End Sub

    Public Overrides Function ProcessOK(ByVal pLine As String) As Boolean
        If m_CanIdle Then
            'If we can idle, then we can proceed with setting up the idle command
            'starting with selecting the correct folder.
            Dim _cmd As New SelectInboxCommand(m_Client)
            m_Client.CmdQueue.Enqueue(_cmd)
        Else
            'Server says we can't idle so we can't continue.
            Throw New Exception("Server does not support IMAP IDLE!")
        End If
        Return MyBase.ProcessOK(pLine)
    End Function

    Public Overrides Function ProcessNO(ByVal pLine As String) As Boolean
        Throw New CantContinueException("Capabilities refused, can't continue", pLine)
    End Function

    Public Overrides Function ProcessBAD(ByVal pLine As String) As Boolean
        Throw New CantContinueException("Capabilities error, can't continue", pLine)
    End Function
End Class

''' <summary>
''' Attempt to logon to the IMAP server.
''' </summary>
''' <remarks></remarks>
Public Class LoginCommand
    Inherits BaseCommand

    Private Const CMD_NAME As String = "LOGIN"

    Private m_UserName As String
    Private m_Password As String

    Public Sub New(ByVal pClient As ImapClient, ByVal pUserName As String, ByVal pPassword As String)
        MyBase.New(pClient)

        m_UserName = pUserName
        m_Password = pPassword
    End Sub

    Public Overrides Sub Trace(ByVal pCmd As String)
        System.Diagnostics.Trace.WriteLine(pCmd.Replace(m_UserName, "MASKED").Replace(m_Password, "MASKED"))
    End Sub

    Public Overrides Sub Send()
        SendCore(CMD_NAME & " " & m_UserName & " " & m_Password)
    End Sub

    Public Overrides Function ProcessOK(ByVal pLine As String) As Boolean
        m_Client.CmdQueue.Enqueue(New CapabilitiesCommand(m_Client))
        System.Diagnostics.Trace.WriteLine("Login ok, proceeding to capabilities.")
        Return MyBase.ProcessOK(pLine)
    End Function

    Public Overrides Function ProcessBAD(ByVal pLine As String) As Boolean
        Throw New CantContinueException("Error on login, please check credentials.", pLine.Replace(m_UserName, "MASKED").Replace(m_Password, "MASKED"))
    End Function

    Public Overrides Function ProcessNO(ByVal pLine As String) As Boolean
        Throw New CantContinueException("Login refused, please check credentials.", pLine.Replace(m_UserName, "MASKED").Replace(m_Password, "MASKED"))
    End Function
End Class

''' <summary>
''' Command to select the correct inbox. This will need to be genericized if anyone needs to select an 
''' inbox that is not called inbox.
''' </summary>
''' <remarks></remarks>
Public Class SelectInboxCommand
    Inherits BaseCommand

    Private Const CMD_NAME As String = "SELECT"
    Private Const INBOX As String = "INBOX"

    Public Sub New(ByVal pClient As ImapClient)
        MyBase.New(pClient)
    End Sub

    Public Overrides Sub Send()
        SendCore(CMD_NAME & " " & INBOX)
    End Sub

    Public Overrides Function ProcessOK(ByVal pLine As String) As Boolean
        m_Client.CmdQueue.Enqueue(New IdleCommand(m_Client))

        Return MyBase.ProcessOK(pLine)
    End Function

    Public Overrides Function ProcessBAD(ByVal pLine As String) As Boolean
        Throw New CantContinueException("Error selecting inbox.", pLine)
    End Function

    Public Overrides Function ProcessNO(ByVal pLine As String) As Boolean
        Throw New CantContinueException("Selecting inbox refused.", pLine)
    End Function
End Class

''' <summary>
''' Used for terminating the IDLE command
''' </summary>
''' <remarks></remarks>
Public Class DoneCommand
    Inherits BaseCommand

    Public Sub New(ByVal pClient As ImapClient)
        MyBase.New(pClient)
    End Sub

    Public Overrides Sub Send()
        Trace("DONE" & vbCrLf)
        m_Client.Send("DONE" & vbCrLf)
    End Sub
End Class

''' <summary>
''' The meat of the application, this command is for telling the server we wan't unsolicited mailbox
''' update notifications. I think its kind of a dumb name for the command to be honest.
''' </summary>
''' <remarks></remarks>
Public Class IdleCommand
    Inherits BaseCommand

    Private Shared m_Lock As New Object
    Private Shared m_Exists As Long = -1

    Public Sub New(ByVal pClient As ImapClient)
        MyBase.New(pClient)
    End Sub

    Private Const CMD_NAME As String = "IDLE"

    Private m_Idling As Boolean = False

    Public Overrides Function ProcessContinuation(ByVal pLine As String) As Boolean
        SyncLock m_Lock
            m_Idling = True
            'when the idle command is successfully set up, we want to raise a new mail
            'notification in case new mail was delivered while our service was not running.
            'Its possible that this should be configurable.
            m_Client.RaiseNotifyMail()
            Return True
        End SyncLock
    End Function

    ''' <summary>
    ''' We need to keep the server from dropping our connection, so this will end this IDLE command,
    ''' and start a new one.
    ''' </summary>
    ''' <remarks></remarks>
    Public Sub Refresh()
        System.Diagnostics.Trace.WriteLine("refreshing idle")
        Dim _noop As New DoneCommand(m_Client)
        _noop.Send()

        m_IsCompleted = True
        Dim _idle As New IdleCommand(m_Client)
        m_Client.CmdQueue.Enqueue(_idle)
    End Sub

    Public Overrides Function ProcessUntagged(ByVal pLine As String) As Boolean
        SyncLock m_Lock
            'The server will send exists commands if the mailbox contents change.
            If pLine.Contains("EXISTS") Then
                System.Diagnostics.Trace.WriteLine("exist arrived")
                'Extract the mailbox count.
                Dim _match As System.Text.RegularExpressions.Match = System.Text.RegularExpressions.Regex.Match(pLine, "\s+(?<num>\d+)\s+EXISTS")
                Dim _num As Long = CLng(_match.Result("${num}"))
                System.Diagnostics.Trace.WriteLine("new mailbox count: " & _num.ToString)
                'If the number of messages has changed since last time, there is new mail.
                If _num <> m_Exists Then
                    m_Client.RaiseNotifyMail()
                End If
                m_Exists = _num
            End If
            Return True
        End SyncLock
    End Function

    Public Overrides Sub Send()
        SyncLock m_Lock
            SendCore(CMD_NAME)
        End SyncLock
    End Sub
End Class

''' <summary>
''' The imap client. Not a full imap client, only serves to set up the IDLE command and listen for 
''' EXISTS notifications.
''' </summary>
''' <remarks></remarks>
Public Class ImapClient
    ''' <summary>
    ''' The IMAP server to connect to.
    ''' </summary>
    ''' <remarks></remarks>
    Private m_Server As String

    ''' <summary>
    ''' The PORT of the IMAP server to connect to.
    ''' </summary>
    ''' <remarks></remarks>
    Private m_Port As Integer

    ''' <summary>
    ''' A data buffer for sending and receiving.
    ''' </summary>
    ''' <remarks></remarks>
    Private m_Buffer As String

    ''' <summary>
    ''' The currently running commands
    ''' </summary>
    ''' <remarks></remarks>
    Private m_RunningCommands As New List(Of ICommand)
    ''' <summary>
    ''' The commands waiting to be started
    ''' </summary>
    ''' <remarks></remarks>
    Private m_CmdQueue As New Generic.Queue(Of ICommand)

    ''' <summary>
    ''' The username to use to authenticate.
    ''' </summary>
    ''' <remarks></remarks>
    Private m_UserName As String

    ''' <summary>
    ''' The password to use to authenticate.
    ''' </summary>
    ''' <remarks></remarks>
    Private m_Password As String

    ''' <summary>
    ''' How often should we refresh the IDLE command.
    ''' </summary>
    ''' <remarks></remarks>
    Private m_Refresh As Integer

    ''' <summary>
    ''' What was the last number of messages to exist.
    ''' </summary>
    ''' <remarks></remarks>
    Private m_Exists As Long

    ''' <summary>
    ''' An event to raise if we think new mail has arrived.
    ''' </summary>
    ''' <remarks></remarks>
    Public Event NotifyMail()

    ''' <summary>
    ''' The commands waiting to be started.
    ''' </summary>
    ''' <value></value>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Public ReadOnly Property CmdQueue() As Generic.Queue(Of ICommand)
        Get
            Return m_CmdQueue
        End Get
    End Property

    Public Sub New(ByVal pServer As String, ByVal pPort As Integer, ByVal pRefresh As Integer)
        m_Server = pServer
        m_Port = pPort
        m_Refresh = pRefresh
    End Sub

    Private m_Socket As Integer

    <DllImport("SSLLib")> _
    Public Shared Function sslsocketclose(ByVal pSocket As Integer, ByRef pError As Integer) As Integer
    End Function

    <DllImport("SSLLib")> _
    Public Shared Function sslsocket(ByVal pHostName As String, ByVal pPort As Integer, ByVal pErrorString As Text.StringBuilder) As Integer
    End Function

    <DllImport("SSLLib")> _
    Public Shared Function sslsend(ByVal pSocket As Integer, ByVal pBuffer As String, ByRef pError As Integer) As Integer
    End Function

    <DllImport("SSLLib")> _
    Public Shared Function sslrecv(ByVal pSocket As Integer, ByVal pBuffer As Text.StringBuilder, ByVal pTimeoutMinutes As Integer, ByRef pError As Integer) As Integer
    End Function

    ''' <summary>
    ''' Read from the socket. Will timeout based on the value of Refresh.
    ''' </summary>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Public Function Read() As String
        Dim _buffer As New Text.StringBuilder(1024)

        Dim _error As Integer = 0
        Dim _err As Integer = sslrecv(m_Socket, _buffer, m_Refresh, _error)
        If _err = 0 Then
            Trace.WriteLine("timeout")
            Throw New TimeoutException()
        End If
        If _err < 0 Then
            Trace.WriteLine("socket error on recv. " & _error.ToString)
            Throw New Exception("socket error on recv. " & _error.ToString)
        End If
        Return _buffer.ToString
    End Function

    ''' <summary>
    ''' Send data on the socket.
    ''' </summary>
    ''' <param name="pBuffer"></param>
    ''' <remarks></remarks>
    Public Sub Send(ByVal pBuffer As String)
        Dim _error As Integer = 0
        Dim _Err As Integer = sslsend(m_Socket, pBuffer, _error)
        If _Err < 0 Then
            Trace.WriteLine("socket error on send. " & _error.ToString)
            Throw New Exception("socket error on send. " & _error.ToString)
        End If
    End Sub

    ''' <summary>
    ''' Set up the socket and initiate the logon sequence
    ''' </summary>
    ''' <remarks></remarks>
    Public Sub SetUpSocket()
        m_Socket = 0
        m_CmdQueue.Clear()

        Dim _sb As New Text.StringBuilder(1024)
        m_Socket = sslsocket(m_Server, m_Port, _sb)

        If m_Socket < 0 Then
            Trace.WriteLine("Error setting up socket: " & _sb.ToString)
            Throw New Exception(_sb.ToString)
        End If

        Logon()
    End Sub

    ''' <summary>
    ''' Send the logon command.
    ''' </summary>
    ''' <remarks></remarks>
    Private Sub Logon()
        Dim _cmd As New LoginCommand(Me, m_UserName, m_Password)
        m_RunningCommands.Add(_cmd)
        _cmd.Send()
    End Sub

    Public Sub ResetIdle()
        Trace.WriteLine("resetting idle")
        Dim _idlecmd As IdleCommand = Nothing
        For Each _cmd As BaseCommand In m_RunningCommands
            If TypeOf _cmd Is IdleCommand Then
                _idlecmd = CType(_cmd, IdleCommand)
            End If
        Next
        If _idlecmd IsNot Nothing Then
            _idlecmd.Refresh()
        End If
        Touch()
    End Sub

    ''' <summary>
    ''' Connect using the provided credentials.
    ''' </summary>
    ''' <param name="pUserName">The username to use to authenticate.</param>
    ''' <param name="pPassword">The password to use to authenticate.</param>
    ''' <remarks></remarks>
    Public Sub Connect(ByVal pUserName As String, ByVal pPassword As String)
        m_UserName = pUserName
        m_Password = pPassword

        SetUpSocket()
        Touch()

        Dim _str As String = Read()
        Dim _lastRetry As Date = DateTime.Now
        Dim _retryTimer As TimeSpan = New TimeSpan(0, 0, 5)
        While True
            Try
                ReadPump(_str)
                If Date.Now.Subtract(m_LastRefresh) > New TimeSpan(0, m_Refresh, 0) Then
                    Trace.WriteLine("Normal refresh has expired.")
                    ResetIdle()
                End If
            Catch _ccex As CantContinueException
                Trace.WriteLine("Got can't continue exception in the read loop, bubbling up.")
                Throw
            Catch ex As Exception
                Trace.WriteLine("Error doing a read, trying to reset up connection. " & ex.Message)

                If DateTime.Now.Subtract(_lastRetry) > New TimeSpan(0, 20, 0) Then
                    Trace.WriteLine("last retry was a while ago, resetting the timespan")
                    _retryTimer = New TimeSpan(0, 0, 5)
                End If
                System.Threading.Thread.Sleep(CInt(_retryTimer.TotalMilliseconds))
                _retryTimer = TimeSpan.FromMilliseconds(_retryTimer.TotalMilliseconds * 2)
                Trace.WriteLine("retry timer is now: " & _retryTimer.TotalMilliseconds)
                _lastRetry = DateTime.Now

                'If we have a problem doing a read, try to reset up the connection.
                SetUpSocket()
            End Try
            Try
                _str = Read()
            Catch ex As TimeoutException
                'We have timed out, we need to reset the idle command, and then continue on.
                Trace.WriteLine("Got timeout exception")
                ResetIdle()
            End Try
        End While
    End Sub

    ''' <summary>
    ''' Process some data we received from the IMAP server.
    ''' </summary>
    ''' <param name="pStr"></param>
    ''' <remarks></remarks>
    Public Sub ReadPump(ByVal pStr As String)
        Trace.WriteLine("read: " & pStr.Replace(m_UserName, "MASKED").Replace(m_Password, "MASKED"))

        m_Buffer += pStr
        While m_Buffer.Contains(vbCrLf)
            Dim _lineBreak As Integer = m_Buffer.IndexOf(vbCrLf)
            Dim _line As String = m_Buffer.Substring(0, _lineBreak)
            m_Buffer = m_Buffer.Substring(_lineBreak + 2, m_Buffer.Length - (_lineBreak + 2))
            ProcessLine(_line)
        End While
    End Sub

    ''' <summary>
    ''' Process a line of data that we received from the IMAP server
    ''' </summary>
    ''' <param name="pLine"></param>
    ''' <remarks></remarks>
    Public Sub ProcessLine(ByVal pLine As String)
        If pLine.StartsWith("* BYE") Then
            Trace.WriteLine("System disconnected unexpectedly.")
            Throw New Exception("System disconnected unexpectedly.")
        End If

        Dim _Remove As New List(Of ICommand)

        'See if any of the running commands can accept the input.
        For Each _comm As ICommand In m_RunningCommands
            If _comm.Process(pLine) Then
                If _comm.IsCompleted Then
                    _Remove.Add(_comm)
                End If
                Exit For
            End If
            'If the current command has become completed, we are done with it.
            If _comm.IsCompleted Then
                _Remove.Add(_comm)
            End If
        Next

        'Remove any commands that have completed.
        For Each _comm As ICommand In _Remove
            m_RunningCommands.Remove(_comm)
        Next

        'Start any commands that are waiting to be started
        While m_CmdQueue.Count > 0
            Dim _cmd As ICommand = m_CmdQueue.Dequeue
            m_RunningCommands.Add(_cmd)
            _cmd.Send()
        End While
    End Sub

    Private m_LastRefresh As Date
    Public Sub Touch()
        m_LastRefresh = Date.Now
    End Sub

    'Do something to disconnect the client?
    Public Sub Disconnect()
        Dim _error As Integer = 0
        Dim _err As Integer = sslsocketclose(m_Socket, _error)
        If _err < 0 Then
            Trace.Write("error disconnecting. " & _error.ToString)
        End If
    End Sub

    ''' <summary>
    ''' Raise the event that indicates we think there is new mail arriving.
    ''' </summary>
    ''' <remarks></remarks>
    Public Sub RaiseNotifyMail()
        RaiseEvent NotifyMail()
    End Sub

End Class